package edu.mcw.rgd.clinminer.umls.indexing;

import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.Stack;

import org.apache.commons.lang.StringUtils;
import org.apache.log4j.Logger;
import org.springframework.context.support.ClassPathXmlApplicationContext;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import edu.mcw.rgd.clinminer.data.ClinminerException;
import edu.mcw.rgd.clinminer.domain.GeneticLab;
import edu.mcw.rgd.clinminer.domain.ImagingPhenotype;
import edu.mcw.rgd.clinminer.domain.Measurement;
import edu.mcw.rgd.clinminer.domain.Medication;
import edu.mcw.rgd.clinminer.domain.NormalPhenotype;
import edu.mcw.rgd.clinminer.domain.Participant;
import edu.mcw.rgd.clinminer.domain.Phenotype;
import edu.mcw.rgd.clinminer.domain.Procedure;
import edu.mcw.rgd.clinminer.domain.Study;
import edu.mcw.rgd.clinminer.domain.common.Measurable;
import edu.mcw.rgd.clinminer.domain.common.Recordable;
import edu.mcw.rgd.clinminer.umls.domain.DataIndex;
import edu.mcw.rgd.clinminer.umls.domain.MrConso;
import edu.mcw.rgd.clinminer.umls.domain.MrHier;
import edu.mcw.rgd.clinminer.umls.type.AUI;
import edu.mcw.rgd.clinminer.umls.type.CUI;

@Component
public class CUIMapper {
	private static Logger log = Logger.getLogger(CUIMapper.class);
	private static String[] header;
	private static Map<Study, Map<CUI, Integer>> studyCUIIndex;
	private static Map<Study, Map<AUI, Integer>> studyAUIIndex;
	private static Map<Study, Set<AUI>> hiddenNodes;
	private static Map<Study, Map<AUI, Set<AUI>>> parentNodes;
	private static Map<AUI, String> AUISources;

	public static void main(String[] args) throws ClinminerException {
		new ClassPathXmlApplicationContext(
				"classpath*:META-INF/spring/applicationContext*.xml");
		// log.debug(MrHier.getAllParentAUIs(new AUI("A16962310")));
		indexerTask();
	}

	@Scheduled(cron = "0 0 0 * * MON-FRI")
	public static void indexerTask() throws ClinminerException {
		log.info("Starting scheduled task");
		mapAnnotationsToCUIs();
		remapCUIsToAUIs();
		propagateInformationContent();
		hideNodes();
		createIndexWithoutHiddenNodes();
		storeIndex();
	}

	private static void storeIndex() {
		log.info("Current index size " + DataIndex.countDataIndexes());
		new DataIndex().clearTable();
		log.info("Index removed | size " + DataIndex.countDataIndexes());

		for (Study s : studyAUIIndex.keySet()) {
			for (Entry<AUI, Set<AUI>> e : parentNodes.get(s).entrySet()) {
				AUI aui = e.getKey();
				for (AUI paui : e.getValue()) {
					DataIndex dataIndex = new DataIndex();
					dataIndex.setAui(aui.toString());
					dataIndex.setPaui(paui.toString());
					dataIndex.setStudy(s);
					dataIndex.setSab(AUISources.get(aui));
					dataIndex.setAnnotationStudyCount(studyAUIIndex.get(s).get(
							aui));
					dataIndex.setAnnotationTotalCount(getTotalCount(aui));
					dataIndex.persist();
				}
				// rootNode has no parents
				if (e.getValue().isEmpty()) {
					DataIndex dataIndex = new DataIndex();
					dataIndex.setAui(aui.toString());
					dataIndex.setStudy(s);
					dataIndex.setSab(AUISources.get(aui));
					dataIndex.setAnnotationStudyCount(studyAUIIndex.get(s).get(
							aui));
					dataIndex.setAnnotationTotalCount(getTotalCount(aui));
					dataIndex.persist();
				}
			}
			// need hidden nodes to display counts if all terms visible
			for (AUI aui : hiddenNodes.get(s)) {
				DataIndex dataIndex = new DataIndex();
				dataIndex.setAui(aui.toString());
				dataIndex.setStudy(s);
				dataIndex.setSab(AUISources.get(aui));
				dataIndex
						.setAnnotationStudyCount(studyAUIIndex.get(s).get(aui));
				dataIndex.setAnnotationTotalCount(getTotalCount(aui));
				dataIndex.persist();
			}
		}
		log.info("Current index size " + DataIndex.countDataIndexes());
	}

	private static Integer getTotalCount(AUI aui) {
		Integer totalCount = 0;
		for (Study s : studyAUIIndex.keySet()) {
			if (studyAUIIndex.get(s).containsKey(aui)) {
				totalCount += studyAUIIndex.get(s).get(aui);
			}
		}
		return totalCount;
	}

	private static void createIndexWithoutHiddenNodes()
			throws ClinminerException {
		parentNodes = new HashMap<Study, Map<AUI, Set<AUI>>>();

		for (Study s : studyAUIIndex.keySet()) {
			parentNodes.put(s, new HashMap<AUI, Set<AUI>>());

			log.info("Creating indexed tree STUDY " + s.getName());

			for (AUI aui : studyAUIIndex.get(s).keySet()) {
				// skip hidden nodes
				if (hiddenNodes.get(s).contains(aui)) {
					continue;
				}

				parentNodes.get(s).put(aui,
						getAllDirectNonHiddenParentAUIs(s, aui));

			}
		}
	}

	public static Set<AUI> getAllDirectNonHiddenParentAUIs(Study s, AUI aui)
			throws ClinminerException {
		Set<AUI> result = new HashSet<AUI>();
		Stack<AUI> stack = new Stack<AUI>();
		// seed the stack
		stack.add(aui);

		while (stack.size() != 0) {
			if (result.contains(stack.peek())) {
				stack.pop();
				continue;
			}
			// process deeper
			if (aui.equals(stack.peek())
					|| hiddenNodes.get(s).contains(stack.peek())) {
				stack.addAll(MrHier.getParentAUIs(stack.pop()));
			} else {
				// until non-hidden parents are found
				result.add(stack.pop());
			}
		}
		result.remove(aui);
		return result;
	}

	private static void hideNodes() throws ClinminerException {
		hiddenNodes = new HashMap<Study, Set<AUI>>();
		for (Study s : studyAUIIndex.keySet()) {
			hiddenNodes.put(s, new HashSet<AUI>());
			log.info("Hiding nodes STUDY " + s.getName() + " COUNT "
					+ studyAUIIndex.get(s).size());
			for (Entry<AUI, Integer> nodeEntry : studyAUIIndex.get(s)
					.entrySet()) {

				Integer nodeIC = nodeEntry.getValue();
				Integer maxChildIC = 0;

				for (AUI child : MrHier.getChildAUIs(nodeEntry.getKey())) {
					// as long as the node contained information downstream
					if (studyAUIIndex.get(s).containsKey(child)) {
						maxChildIC = Math.max(maxChildIC, studyAUIIndex.get(s)
								.get(child));
					}
				}
				// hide node if it doesn't bring more information
				if (nodeIC <= maxChildIC) {
					hiddenNodes.get(s).add(nodeEntry.getKey());
				}

			}
			log.info("\t" + hiddenNodes.get(s).size() + " nodes hidden");
		}

	}

	// note: this will add parent nodes previously unmapped into the map
	private static void propagateInformationContent() throws ClinminerException {
		for (Study s : studyAUIIndex.keySet()) {
			log.info("Propagating weights STUDY " + s.getName() + " COUNT "
					+ studyAUIIndex.get(s).size());

			// local copy as modifying on the fly
			for (Entry<AUI, Integer> nodeEntry : new HashSet<Entry<AUI, Integer>>(
					studyAUIIndex.get(s).entrySet())) {
				Integer currentNodeIC = nodeEntry.getValue();
				for (AUI parent : MrHier.getAllParentAUIs(nodeEntry.getKey())) {
					Integer count = studyAUIIndex.get(s).get(parent);
					if (count == null) {
						count = 0;
					}
					studyAUIIndex.get(s).put(parent, count + currentNodeIC);
					AUISources.put(parent, AUISources.get(nodeEntry.getKey()));
				}
			}

		}

	}

	// Limits indexed sources for debugging
	private static Collection<String> listUMLSHierarchicalSources() {
		// return Arrays.asList("SNOMEDCT");
		return MrHier.findDistinctSources();
	}

	private static void remapCUIsToAUIs() throws ClinminerException {
		studyAUIIndex = new HashMap<Study, Map<AUI, Integer>>();
		AUISources = new HashMap<AUI, String>();

		for (Study s : studyCUIIndex.keySet()) {
			studyAUIIndex.put(s, new HashMap<AUI, Integer>());

			for (String source : listUMLSHierarchicalSources()) {
				log.info("Remapping STUDY " + s.getName() + " SOURCE " + source);
				Map<AUI, Integer> sourceIndex = new HashMap<AUI, Integer>();
				Map<CUI, Integer> index = studyCUIIndex.get(s);

				// remap back to source
				for (Entry<CUI, Integer> e : index.entrySet()) {
					Collection<AUI> auis = MrHier.findAUIbyCUIandSAB(
							e.getKey(), source, false);
					if (auis.isEmpty()) {
						auis = MrHier.findAUIbyCUIandSAB(e.getKey(), source,
								true);
						if (!auis.isEmpty()) {
							log.warn("All AUIs suppressed for CUI "
									+ e.getKey() + " (" + e.getValue() + ")");
						}
					}
					// cui might not have been represented in the source
					// but if not empty with supressed terms included
					// original mapping was made to an obsolete term

					for (AUI aui : auis) {
						Integer count = sourceIndex.get(aui);
						if (count == null) {
							count = 0;
						}
						sourceIndex.put(aui, count + e.getValue());
						AUISources.put(aui, source);
					}
				}

				studyAUIIndex.get(s).putAll(sourceIndex);
			}
		}
	}

	private static void mapAnnotationsToCUIs() throws ClinminerException {
		studyCUIIndex = new HashMap<Study, Map<CUI, Integer>>();

		Integer countMappings = 0;
		for (Study s : Study.findAllStudies()) {
			for (Recordable r : getAllAnnotations(s)) {
				for (Entry<String, String> e : getOntologyEntries(r)) {
					countMappings++;
					if (e.getValue() != null && e.getKey() == null) {
						// provided term label but no id
						// log.warn(e + " missing id in " + getParticipant(r));
						continue;
					} else if (e.getKey() == null) {
						// otherwise missing id+labels are ok in some parts of
						// the
						// model
						continue;
					}
					List<String> cuis = MrConso.findCuisForCode(e.getKey()
							.trim());
					if (cuis.isEmpty()) {
						cuis = MrConso.findSuppressedCuisForCode(e.getKey()
								.trim());
						if (cuis.isEmpty()) {
							System.out.println(getParticipant(r)
									+ "\tNOT MAPPABLE\t" + e);
						} else {
							System.out.println(getParticipant(r)
									+ "\tSUPRESSED CONCEPT\t" + e + "\t"
									+ toStringCuis(cuis));
						}
					} else if (cuis.size() > 1) {
						System.out.println(getParticipant(r)
								+ "\tMULTIPLE MAPPINGS\t" + e + "\t"
								+ toStringCuis(cuis));
					}

					if (cuis.contains("C1285556")) {
						throw new UnsupportedOperationException(e.getKey());
					}

					for (String cui : cuis) {
						Map<CUI, Integer> mappedCuis = studyCUIIndex.get(s);
						if (mappedCuis == null) {
							mappedCuis = new HashMap<CUI, Integer>();
							studyCUIIndex.put(s, mappedCuis);
						}
						Integer count = mappedCuis.get(cui);
						if (count == null) {
							count = 0;
						}
						count++;
						mappedCuis.put(new CUI(cui), count);
					}
				}
			}
		}
		// report
		// System.out.println("");
		// Integer counter = 0;
		// for (Entry<Study, Map<String, Integer>> e : studyIndex.entrySet()) {
		// log.info("\t" + e.getKey().getName());
		// for (Entry<String, Integer> e2 : e.getValue().entrySet()) {
		// log.info(e.getKey().getName() + "\t" + e2.getKey() + "\t"
		// + e2.getValue());
		// }
		// }

	}

	private static String toStringCuis(List<String> cuis) {
		List<String> entries = new ArrayList<String>();
		for (String cui : cuis) {
			entries.add(cui + "=" + MrConso.findPrefLabelForCui(cui));
		}
		return StringUtils.join(entries, ", ");
	}

	private static List<Entry<String, String>> getOntologyEntries(Recordable r) {
		List<Entry<String, String>> list = new ArrayList<Entry<String, String>>();

		list.add(new AbstractMap.SimpleEntry<String, String>(r
				.getOntologyTermID(), r.getOntologyTermLabel()));

		if (r instanceof Measurable) {
			String id = ((Measurable) r).getLOINCID();
			String label = ((Measurable) r).getLOINCLabel();
			list.add(new AbstractMap.SimpleEntry<String, String>(id, label));
		}

		if (r instanceof GeneticLab) {
			String id = ((GeneticLab) r).getLOINCID();
			String label = ((GeneticLab) r).getLOINCLabel();
			list.add(new AbstractMap.SimpleEntry<String, String>(id, label));
		}

		return list;
	}

	// no need to cache as JPA is already doing it
	private static List<Recordable> getAllAnnotations(Study s) {
		List<Recordable> annotations = new ArrayList<Recordable>();
		for (Participant p : s.getParticipants()) {
			annotations.addAll(p.getPhenotypes());
			annotations.addAll(p.getMedications());
			annotations.addAll(p.getImagingPhenotypes());
			annotations.addAll(p.getMeasurements());
			annotations.addAll(p.getGeneticLabs());
			annotations.addAll(p.getProcedures());
			annotations.addAll(p.getNormalPhenotypes());
		}
		return annotations;
	}

	private static Participant getParticipant(Recordable r) {
		if (r instanceof Phenotype) {
			return ((Phenotype) r).getParticipant();
		}
		if (r instanceof Medication) {
			return ((Medication) r).getParticipant();
		}
		if (r instanceof ImagingPhenotype) {
			return ((ImagingPhenotype) r).getParticipant();
		}
		if (r instanceof Measurement) {
			return ((Measurement) r).getParticipant();
		}
		if (r instanceof GeneticLab) {
			return ((GeneticLab) r).getParticipant();
		}
		if (r instanceof Procedure) {
			return ((Procedure) r).getParticipant();
		}
		if (r instanceof NormalPhenotype) {
			return ((NormalPhenotype) r).getParticipant();
		}
		throw new UnsupportedOperationException();
	}
}